package viterbi;

import java.io.BufferedReader;
import java.io.File;
import java.io.FileReader;
import java.io.IOException;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.LinkedList;


import common.CorpusProcessorNoisy;
import common.Probability;

public class ViterbiAccNew {
	// From training set
	private static HashMap<String, Integer> tCount;
	private static int tTotalCount;
	private static HashMap<String, HashMap<String, Integer>> ttCount;
	private static HashMap<String, HashMap<String, Integer>> wtCount;
	private static HashMap<String, String> mostFrequentSubsequentPOS;
	// Temporary variables
	private static int currentBackPointer;
	private static int currentMaxIndexLastColumn;
	private static ArrayList<ArrayList<String>> currentSentencePosGraph;
	private static ArrayList<ArrayList<Probability>> currentSentencePosProbabilityGraph;
	private static ArrayList<ArrayList<Integer>> currentSentencePosBackPointerGraph;
	private static ArrayList<String> currentSentence;
	private static Probability currentMaxProb;
	private static int[] correctTags;

	public static void main(String[] args) {
		System.out.println("Applying the Viterbi tagger to the development set...");
		tCount = new HashMap<String, Integer>();
		ttCount = new HashMap<String, HashMap<String, Integer>>();
		wtCount = new HashMap<String, HashMap<String, Integer>>();
		mostFrequentSubsequentPOS = new HashMap<String, String>();
		currentBackPointer = 0;
		currentMaxProb = new Probability();
		correctTags = new int[2];

		tTotalCount = CorpusProcessorNoisy.run(tCount, ttCount, wtCount,
				mostFrequentSubsequentPOS);

		try {
			File file = new File("CoNLL2009-ST-English-development-pos.txt");
			currentSentence = new ArrayList<String>();
			ArrayList<String> currentSentencePOS = new ArrayList<String>();
			BufferedReader br = new BufferedReader(new FileReader(file));
			String line = null;

			while ((line = br.readLine()) != null) {
				if (!line.isEmpty()) {
					String[] cols = line.split("\t");
						currentSentence.add(cols[2]);
						currentSentencePOS.add(cols[4]);
				} else {
						LinkedList<String> taggedPOS = tag();
						for (int i = 0; i < currentSentencePOS.size(); i++) {
							if (currentSentencePOS.get(i).equals(
									taggedPOS.get(i))) {
								correctTags[0]++;
							} else {
								correctTags[1]++;
							}
						}
						currentSentence = new ArrayList<String>();
						currentSentencePOS = new ArrayList<String>();
				}

			}
			double acc = (double) correctTags[0]
					/ (correctTags[0] + correctTags[1]);
			System.out.println("Done tagging. Accuracy was " + acc);
		} catch (IOException e) {
			e.printStackTrace();
		}
	}

	private static LinkedList<String> tag() {
		currentSentencePosProbabilityGraph = new ArrayList<ArrayList<Probability>>();
		currentSentencePosBackPointerGraph = new ArrayList<ArrayList<Integer>>();
		currentSentencePosGraph = getPosGraph(currentSentence);
		int[] possibilities = new int[currentSentence.size()];
		for (int i = 0; i < currentSentencePosGraph.size(); i++) {
			possibilities[i] = currentSentencePosGraph.get(i).size();
			if (i == 0) {
				fillFirstColumn();
			} else {
				Probability tempMaxProb = new Probability();
				ArrayList<Probability> probs = new ArrayList<Probability>();
				ArrayList<Integer> backPointers = new ArrayList<Integer>();
				for (int j = 0; j < possibilities[i]; j++) {
					for (int k = 0; k < possibilities[i - 1]; k++) {
						Probability p = getColumnProb(i, j, k);
						if (p.compareTo(tempMaxProb) == 1) {
							tempMaxProb = p;
							currentBackPointer = k;
						}
					}
					tempMaxProb = tempMaxProb.multiply(getWGivenTProb(i, j));
					probs.add((Probability) tempMaxProb.clone());
					backPointers.add(currentBackPointer);
					tempMaxProb.setProbability(0);
				}
				currentSentencePosProbabilityGraph.add(probs);
				currentSentencePosBackPointerGraph.add(backPointers);
			}
		}

		currentMaxProb = new Probability();
		for (int i = 0; i < currentSentencePosProbabilityGraph.get(
				currentSentencePosProbabilityGraph.size() - 1).size(); i++) {
			Probability p = currentSentencePosProbabilityGraph.get(
					currentSentencePosProbabilityGraph.size() - 1).get(i);
			if (p.compareTo(currentMaxProb) == 1) {
				currentMaxProb = p;
				currentMaxIndexLastColumn = i;
			}
		}

		LinkedList<String> POStags = backTrack();

		return POStags;
	}

	private static ArrayList<ArrayList<String>> getPosGraph(
			ArrayList<String> words) {
		ArrayList<ArrayList<String>> result = new ArrayList<ArrayList<String>>();
		for (int i = 0; i < words.size(); i++) {
			ArrayList<String> a = new ArrayList<String>();
			if (wtCount.containsKey(words.get(i))) {
				a = new ArrayList<String>(wtCount.get(words.get(i)).keySet());
			} else if (i == 0) {
				a.add(mostFrequentSubsequentPOS.get("<s>"));
			} else {
				for (int j = 0; j < result.get(i - 1).size(); j++) {
					a.add(mostFrequentSubsequentPOS.get(result.get(i - 1)
							.get(j)));
				}
			}
			result.add(a);
		}
		return result;
	}

	private static void fillFirstColumn() {
		ArrayList<Probability> probs = new ArrayList<Probability>();
		for (int i = 0; i < currentSentencePosGraph.get(0).size(); i++) {
			Probability prob = new Probability((double) 1);
			String prevPOS = "<s>";
			int prevPOSCount = tCount.get(prevPOS);
			// P(ti | ti-1) term
			String currentPOS = currentSentencePosGraph.get(0).get(i);
			int currentPOSCount = tCount.get(currentPOS);
			Integer lastTwoPOSCount = ttCount.get(prevPOS).get(currentPOS);
			if (lastTwoPOSCount != null) {
				prob = prob.multiply(new Probability((double) lastTwoPOSCount
						/ prevPOSCount));
			} else {
				// Back-off to unigram approximation
				prob = prob.multiply(new Probability((double) currentPOSCount
						/ tTotalCount));
				prob = prob.multiply(new Probability((double) prevPOSCount
						/ tTotalCount));
			}

			// P(wi | ti) term
			if (wtCount.containsKey(currentSentence.get(0))) {
				Integer currentWordPOScount = wtCount.get(
						currentSentence.get(0)).get(currentPOS);
				if (currentWordPOScount != null) {
					prob = prob.multiply(new Probability(
							(double) currentWordPOScount / currentPOSCount));
				} else {
					// Back-off to unigrams
					prob = prob.multiply(new Probability(
							(double) currentPOSCount / tTotalCount));
				}
			} else {
				// Back-off to unigrams
				prob = prob.multiply(new Probability((double) currentPOSCount
						/ tTotalCount));
			}
			probs.add(prob);
		}
		currentSentencePosProbabilityGraph.add(probs);
	}

	private static Probability getColumnProb(int i, int j, int k) {
		Probability prob = new Probability((double) 1);
		String prevPOS = currentSentencePosGraph.get(i - 1).get(k);
		int prevPOSCount = tCount.get(prevPOS);
		// P(ti | ti-1) term
		String currentPOS = currentSentencePosGraph.get(i).get(j);
		int currentPOSCount = tCount.get(currentPOS);
		Integer lastTwoPOSCount = ttCount.get(prevPOS).get(currentPOS);
		if (lastTwoPOSCount != null) {
			prob = prob.multiply(new Probability((double) lastTwoPOSCount
					/ prevPOSCount));
		} else {
			// Back-off to unigram approximation
			prob = prob.multiply(new Probability((double) currentPOSCount
					/ tTotalCount));
			prob = prob.multiply(new Probability((double) prevPOSCount
					/ tTotalCount));
		}

		// Multiply by previous max probability
		prob = prob.multiply(currentSentencePosProbabilityGraph.get(i - 1).get(
				k));

		return prob;
	}

	private static Probability getWGivenTProb(int i, int j) {
		Probability prob = new Probability((double) 1);
		String currentPOS = currentSentencePosGraph.get(i).get(j);
		int currentPOSCount = tCount.get(currentPOS);
		// P(wi | ti) term
		if (wtCount.containsKey(currentSentence.get(i))) {
			Integer currentWordPOScount = wtCount.get(currentSentence.get(i))
					.get(currentPOS);
			if (currentWordPOScount != null) {
				prob = prob.multiply(new Probability(
						(double) currentWordPOScount / currentPOSCount));
			} else {
				// Back-off to unigrams
				prob = prob.multiply(new Probability((double) currentPOSCount
						/ tTotalCount));
			}
		} else {
			// Back-off to unigrams
			prob = prob.multiply(new Probability((double) currentPOSCount
					/ tTotalCount));
		}

		return prob;
	}

	private static LinkedList<String> backTrack() {
		if (currentSentencePosGraph.size()==1){
			LinkedList<String> path = new LinkedList<String>();
			path.add(currentSentencePosGraph.get(0).get(currentMaxIndexLastColumn));
			return path;
		}
		LinkedList<String> path = new LinkedList<String>();
		path.add(currentSentencePosGraph
				.get(currentSentencePosGraph.size() - 1).get(
						currentMaxIndexLastColumn));
		int nextIndex = currentSentencePosBackPointerGraph.get(
				currentSentencePosBackPointerGraph.size() - 1).get(
				currentMaxIndexLastColumn);
		for (int i = currentSentencePosBackPointerGraph.size() - 1; i > 0; i--) {
			path.addFirst(currentSentencePosGraph.get(i).get(nextIndex));
			nextIndex = currentSentencePosBackPointerGraph.get(i - 1).get(
					nextIndex);
		}
		path.addFirst(currentSentencePosGraph.get(0).get(nextIndex));

		return path;
	}
}